#include "SE3GaussianPDF.h"
#include <TooN/TooN.h>
#include <mrpt/base.h>

using namespace mrpt::utils;
using namespace mrpt::poses;

SE3GaussianPDF::SE3GaussianPDF()
{
  se3Mean = TooN::SE3<>();
  m66Covariance = TooN::Zeros;
  bHasRotation = false;
  bHasTranslation = false;
}


SE3GaussianPDF::SE3GaussianPDF(const TooN::SE3<> &se3RT)
{
  se3Mean = se3RT;
  m66Covariance = TooN::Zeros;
  bHasRotation = true;
  bHasTranslation = true;
}


SE3GaussianPDF::SE3GaussianPDF(const TooN::SE3<> &se3RT, const TooN::Matrix<6,6> &m66CovRT)
{
  se3Mean = se3RT;
  m66Covariance = m66CovRT;
  bHasRotation = true;
  bHasTranslation = true;
}


SE3GaussianPDF::SE3GaussianPDF(const TooN::SO3<> &so3R, const TooN::Vector<3> &v3CovR, const TooN::Vector<3> &v3T, const TooN::Vector<3> &v3CovT)
{
  se3Mean = TooN::SE3<>(so3R, v3T);
  m66Covariance = TooN::Zeros;
  for(int i = 0; i < 3; ++i)
    m66Covariance(i,i) = v3CovT[i];
  for(int i = 0; i < 3; ++i)
    m66Covariance(i+3,i+3) = v3CovR[i];
  bHasRotation = true;
  bHasTranslation = true;
}


SE3GaussianPDF::SE3GaussianPDF(const TooN::SO3<> &so3R, const TooN::Vector<3> &v3T, const TooN::Matrix<6,6> &m66CovRT)
{
  se3Mean = TooN::SE3<>(so3R, v3T);
  m66Covariance = m66CovRT;
  bHasRotation = true;
  bHasTranslation = true;
}


SE3GaussianPDF::SE3GaussianPDF(const TooN::SO3<> &so3R, const TooN::Vector<3> &v3CovR)
{
  se3Mean.get_translation() = TooN::Zeros;
  se3Mean.get_rotation() = so3R;
  m66Covariance = TooN::Zeros;
  for(int i = 0; i < 3; ++i)
    m66Covariance(i+3,i+3) = v3CovR[i];
  bHasRotation = true;
  bHasTranslation = false;
}


SE3GaussianPDF::SE3GaussianPDF(const TooN::SO3<> &so3R, const TooN::Matrix<3,3> &m33CovR)
{
  se3Mean.get_translation() = TooN::Zeros;
  se3Mean.get_rotation() = so3R;
  m66Covariance = TooN::Zeros;
  for(int i = 0; i < 3; ++i)
    for(int j = 0; j < 3; ++j)
      m66Covariance(i+3,j+3) = m33CovR(i,j);
  bHasRotation = true;
  bHasTranslation = false;
}


void SE3GaussianPDF::operator += (const SE3GaussianPDF &rhs)
{
  if(!bHasRotation || !rhs.bHasRotation) // Need rotation for both to have any useful result
  {
    *this = SE3GaussianPDF(); 
  }
  else if(bHasTranslation && rhs.bHasTranslation) // Have both rotation and translation both so get full pose pdf
  {
    const TooN::Matrix<3,3>& rLocalLHS = se3Mean.get_rotation().get_matrix();
    const TooN::Vector<3>& tLocalLHS = se3Mean.get_translation();
    const TooN::Matrix<3,3>& rLocalRHS = rhs.se3Mean.get_rotation().get_matrix();
    const TooN::Vector<3>& tLocalRHS = rhs.se3Mean.get_translation();

    // Conversion TooN -> MRPT
    CMatrixDouble33 rLHS, rRHS;
    CArrayDouble<3> tLHS, tRHS;
    for(int i = 0; i < 3; ++i)
    {
      tLHS[i] = tLocalLHS[i];
      tRHS[i] = tLocalRHS[i];
      for(int j = 0; j < 3; ++j)
      {
        rLHS(i,j) = rLocalLHS(i,j);
        rRHS(i,j) = rLocalRHS(i,j);
      }
    }

    CMatrixDouble66 covLHS, covRHS;
    for(int i = 0; i < 6; ++i)
    {
      for(int j = 0; j < 6; ++j)
      {
        covLHS(i,j) = m66Covariance(i,j);
        covRHS(i,j) = rhs.m66Covariance(i,j);
      }
    }

    CPose3D poseLHS(rLHS, tLHS);
    CPose3D poseRHS(rRHS, tRHS);
    CPose3DPDFGaussian pdfLHS(poseLHS, covLHS);
    CPose3DPDFGaussian pdfRHS(poseRHS, covRHS);

    // Actual pose composition
    pdfLHS += pdfRHS;

    // Conversion MRPT -> TooN
    pdfLHS.getCovarianceAndMean(covLHS, poseLHS);
    TooN::Matrix<3,3> rResult;
    TooN::Vector<3> tResult;
    for(int i = 0; i < 3; ++i)
    {
      tResult[i] = poseLHS.m_coords[i];
      for(int j = 0; j < 3; ++j)
        rResult(i,j) = poseLHS.m_ROT(i,j);
    }
    se3Mean = TooN::SE3<>(TooN::SO3<>(rResult), tResult);

    for(int i = 0; i < 6; ++i)
      for(int j = 0; j < 6; ++j)
        m66Covariance(i,j) = covLHS(i,j);
  }
  else // Only have rotation for both, so just get orientation pdf
  {
    const TooN::Matrix<3,3>& rLocalLHS = se3Mean.get_rotation().get_matrix();
    const TooN::Matrix<3,3>& rLocalRHS = rhs.se3Mean.get_rotation().get_matrix();

    // Conversion TooN -> MRPT
    CMatrixDouble33 rLHS, rRHS;
    CArrayDouble<3> tZero;
    for(int i = 0; i < 3; ++i)
    {
      tZero[i] = 0.0;
      for(int j = 0; j < 3; ++j)
      {
        rLHS(i,j) = rLocalLHS(i,j);
        rRHS(i,j) = rLocalRHS(i,j);
      }
    }

    CMatrixDouble66 covLHS, covRHS;
    covLHS.zeros();
    covRHS.zeros();
    for(int i = 3; i < 6; ++i) //Leave translation part and associated covariances all zeros
    {
      for(int j = 3; j < 6; ++j) //Leave translation part and associated covariances all zeros
      {
        covLHS(i,j) = m66Covariance(i,j);
        covRHS(i,j) = rhs.m66Covariance(i,j);
      }
    }

    CPose3D poseLHS(rLHS, tZero);
    CPose3D poseRHS(rRHS, tZero);
    CPose3DPDFGaussian pdfLHS(poseLHS, covLHS);
    CPose3DPDFGaussian pdfRHS(poseRHS, covRHS);

    // Actual pose composition
    pdfLHS += pdfRHS;

    // Conversion MRPT -> TooN
    pdfLHS.getCovarianceAndMean(covLHS, poseLHS);
    TooN::Matrix<3,3> rResult;
    TooN::Vector<3> tResult;
    for(int i = 0; i < 3; ++i)
    {
      tResult[i] = poseLHS.m_coords[i];
      for(int j = 0; j < 3; ++j)
        rResult(i,j) = poseLHS.m_ROT(i,j);
    }
    se3Mean = TooN::SE3<>(TooN::SO3<>(rResult), tResult);

    for(int i = 0; i < 6; ++i)
      for(int j = 0; j < 6; ++j)
        m66Covariance(i,j) = covLHS(i,j);    

    bHasTranslation = false; // May or may not be true already
  }
}


SE3GaussianPDF SE3GaussianPDF::inverse()
{
  if(!bHasRotation) // Need rotation to have any useful result
  {
    return(SE3GaussianPDF()); 
  }
  else if(bHasTranslation) // Have both rotation and translation so get full pose pdf
  {
    const TooN::Matrix<3,3>& rLocalThis = se3Mean.get_rotation().get_matrix();
    const TooN::Vector<3>& tLocalThis = se3Mean.get_translation();

    // Conversion TooN -> MRPT
    CMatrixDouble33 rThis;
    CArrayDouble<3> tThis;
    for(int i = 0; i < 3; ++i)
    {
      tThis[i] = tLocalThis[i];
      for(int j = 0; j < 3; ++j)
        rThis(i,j) = rLocalThis(i,j);
    }

    CMatrixDouble66 covThis, covInv;
    for(int i = 0; i < 6; ++i)
      for(int j = 0; j < 6; ++j)
        covThis(i,j) = m66Covariance(i,j);

    CPose3D poseThis(rThis, tThis), poseInv;
    CPose3DPDFGaussian pdfThis(poseThis, covThis);
    CPose3DPDFGaussian pdfInv;

    // Actual pose composition
    pdfThis.inverse(pdfInv);

    // Conversion MRPT -> TooN
    pdfInv.getCovarianceAndMean(covInv, poseInv);
    TooN::Matrix<3,3> rInv;
    TooN::Vector<3> tInv;
    for(int i = 0; i < 3; ++i)
    {
      tInv[i] = poseInv.m_coords[i];
      for(int j = 0; j < 3; ++j)
        rInv(i,j) = poseInv.m_ROT(i,j);
    }

    TooN::Matrix<6,6> m66CovarianceInv;
    for(int i = 0; i < 6; ++i)
      for(int j = 0; j < 6; ++j)
        m66CovarianceInv(i,j) = covInv(i,j);

    return(SE3GaussianPDF(TooN::SE3<>(TooN::SO3<>(rInv), tInv), m66CovarianceInv)); 
  }
  else // Only have rotation, so just get orientation pdf
  {
    const TooN::Matrix<3,3>& rLocalThis = se3Mean.get_rotation().get_matrix();

    // Conversion TooN -> MRPT
    CMatrixDouble33 rThis;
    CArrayDouble<3> tZero;
    for(int i = 0; i < 3; ++i)
    {
      tZero[i] = 0.0;
      for(int j = 0; j < 3; ++j)
        rThis(i,j) = rLocalThis(i,j);
    }

    CMatrixDouble66 covThis, covInv;
    covThis.zeros();
    for(int i = 3; i < 6; ++i)
      for(int j = 3; j < 6; ++j)
        covThis(i,j) = m66Covariance(i,j);

    CPose3D poseThis(rThis, tZero), poseInv;
    CPose3DPDFGaussian pdfThis(poseThis, covThis);
    CPose3DPDFGaussian pdfInv;

    // Actual pose composition
    pdfThis.inverse(pdfInv);

    // Conversion MRPT -> TooN
    pdfInv.getCovarianceAndMean(covInv, poseInv);
    TooN::Matrix<3,3> rInv;
    for(int i = 0; i < 3; ++i)
      for(int j = 0; j < 3; ++j)
        rInv(i,j) = poseInv.m_ROT(i,j);

    TooN::Matrix<3,3> m33CovarianceInv;
    for(int i = 0; i < 3; ++i)
      for(int j = 0; j < 3; ++j)
        m33CovarianceInv(i,j) = covInv(i+3,j+3);

    return(SE3GaussianPDF(TooN::SO3<>(rInv), m33CovarianceInv)); 
  }
}


void SE3GaussianPDF::operator -= (const SE3GaussianPDF &rhs)
{
  if(!bHasRotation || !rhs.bHasRotation) // Need rotation for both to have any useful result
  {
    *this = SE3GaussianPDF(); 
  }
  else if(bHasTranslation && rhs.bHasTranslation) // Have both rotation and translation both so get full pose pdf
  {
    const TooN::Matrix<3,3>& rLocalLHS = se3Mean.get_rotation().get_matrix();
    const TooN::Vector<3>& tLocalLHS = se3Mean.get_translation();
    const TooN::Matrix<3,3>& rLocalRHS = rhs.se3Mean.get_rotation().get_matrix();
    const TooN::Vector<3>& tLocalRHS = rhs.se3Mean.get_translation();

    // Conversion TooN -> MRPT
    CMatrixDouble33 rLHS, rRHS;
    CArrayDouble<3> tLHS, tRHS;
    for(int i = 0; i < 3; ++i)
    {
      tLHS[i] = tLocalLHS[i];
      tRHS[i] = tLocalRHS[i];
      for(int j = 0; j < 3; ++j)
      {
        rLHS(i,j) = rLocalLHS(i,j);
        rRHS(i,j) = rLocalRHS(i,j);
      }
    }

    CMatrixDouble66 covLHS, covRHS;
    for(int i = 0; i < 6; ++i)
    {
      for(int j = 0; j < 6; ++j)
      {
        covLHS(i,j) = m66Covariance(i,j);
        covRHS(i,j) = rhs.m66Covariance(i,j);
      }
    }

    CPose3D poseLHS(rLHS, tLHS);
    CPose3D poseRHS(rRHS, tRHS);
    CPose3DPDFGaussian pdfLHS(poseLHS, covLHS);
    CPose3DPDFGaussian pdfRHS(poseRHS, covRHS);

    // Actual pose composition
    pdfLHS -= pdfRHS;

    // Conversion MRPT -> TooN
    pdfLHS.getCovarianceAndMean(covLHS, poseLHS);
    TooN::Matrix<3,3> rResult;
    TooN::Vector<3> tResult;
    for(int i = 0; i < 3; ++i)
    {
      tResult[i] = poseLHS.m_coords[i];
      for(int j = 0; j < 3; ++j)
        rResult(i,j) = poseLHS.m_ROT(i,j);
    }
    se3Mean = TooN::SE3<>(TooN::SO3<>(rResult), tResult);

    for(int i = 0; i < 6; ++i)
      for(int j = 0; j < 6; ++j)
        m66Covariance(i,j) = covLHS(i,j);
  }
  else // Only have rotation for both, so just get orientation pdf
  {
    const TooN::Matrix<3,3>& rLocalLHS = se3Mean.get_rotation().get_matrix();
    const TooN::Matrix<3,3>& rLocalRHS = rhs.se3Mean.get_rotation().get_matrix();

    // Conversion TooN -> MRPT
    CMatrixDouble33 rLHS, rRHS;
    CArrayDouble<3> tZero;
    for(int i = 0; i < 3; ++i)
    {
      tZero[i] = 0.0;
      for(int j = 0; j < 3; ++j)
      {
        rLHS(i,j) = rLocalLHS(i,j);
        rRHS(i,j) = rLocalRHS(i,j);
      }
    }

    CMatrixDouble66 covLHS, covRHS;
    covLHS.zeros();
    covRHS.zeros();
    for(int i = 3; i < 6; ++i) //Leave translation part and associated covariances all zeros
    {
      for(int j = 3; j < 6; ++j) //Leave translation part and associated covariances all zeros
      {
        covLHS(i,j) = m66Covariance(i,j);
        covRHS(i,j) = rhs.m66Covariance(i,j);
      }
    }

    CPose3D poseLHS(rLHS, tZero);
    CPose3D poseRHS(rRHS, tZero);
    CPose3DPDFGaussian pdfLHS(poseLHS, covLHS);
    CPose3DPDFGaussian pdfRHS(poseRHS, covRHS);

    // Actual pose composition
    pdfLHS -= pdfRHS;

    // Conversion MRPT -> TooN
    pdfLHS.getCovarianceAndMean(covLHS, poseLHS);
    TooN::Matrix<3,3> rResult;
    TooN::Vector<3> tResult;
    for(int i = 0; i < 3; ++i)
    {
      tResult[i] = poseLHS.m_coords[i];
      for(int j = 0; j < 3; ++j)
        rResult(i,j) = poseLHS.m_ROT(i,j);
    }
    se3Mean = TooN::SE3<>(TooN::SO3<>(rResult), tResult);

    for(int i = 0; i < 6; ++i)
      for(int j = 0; j < 6; ++j)
        m66Covariance(i,j) = covLHS(i,j);    

    bHasTranslation = false; // May or may not be true already
  }
}


SE3GaussianPDF SE3GaussianPDF::operator + (const SE3GaussianPDF &rhs) const
{
  SE3GaussianPDF result = *this;
  result += rhs;
  return result;
}


SE3GaussianPDF SE3GaussianPDF::operator - (const SE3GaussianPDF &rhs) const
{
  SE3GaussianPDF result = *this;
  result -= rhs;
  return result;
}

